import torch
import scipy.io as scio
from torch.utils.data.distributed import DistributedSampler

device = torch.device('cuda')


# voxel_length = 36
# case = 1000

# make dataset
class DataSet(torch.utils.data.Dataset):

    def __init__(self, voxels, hard_element_stiffness, hard_element_force, soft_element_stiffness, soft_element_force, C_hard, C_soft, jacobi, voxel_length = 36, case = 1000):
        self.voxels = torch.tensor(voxels).float().to(device)
        self.hard_element_stiffness = torch.tensor(hard_element_stiffness).float().to(device)
        self.hard_element_force = torch.tensor(hard_element_force).float().to(device)
        self.soft_element_stiffness = torch.tensor(soft_element_stiffness).float().to(device)
        self.soft_element_force = torch.tensor(soft_element_force).float().to(device)
        self.C_hard = torch.tensor(C_hard).float().to(device)
        self.C_soft = torch.tensor(C_soft).float().to(device)
        self.jacobi = torch.tensor(jacobi).float().to(device)
        self.voxel_length = voxel_length
        self.case = case

        print(hard_element_force.shape)
        # print(jacobi.shape)
        # print(voxels.shape)
        self.voxel_length = voxel_length

    def __getitem__(self, index):
        voxel = self.voxels[index // self.case]
        hard_stiffness = self.hard_element_stiffness[index]
        hard_force = self.hard_element_force[index]
        soft_stiffness = self.soft_element_stiffness[index]
        soft_force = self.soft_element_force[index]
        c_hard = self.C_hard[index]
        c_soft = self.C_soft[index]
        jacobi_ = self.jacobi[index]
        input = (c_hard.reshape(36, 1) @ voxel.reshape(1, -1) + c_soft.reshape(36, 1) @ (1 - voxel).reshape(1, -1)).reshape(36, self.voxel_length,
                                                                                          self.voxel_length, self.voxel_length)
        # input = voxel.contiguous().view(1, -1).contiguous().view(1, voxel_length, voxel_length, voxel_length)

        return input, voxel, hard_stiffness, hard_force, soft_stiffness, soft_force, c_hard, c_soft, jacobi_, index

    def __len__(self):
        return len(self.jacobi)

def getDataLoader_(base_path, data_type, voxel_length, case, batch_size = 8):


    matr1 = scio.loadmat(base_path + 'voxel_matrix.mat')
    voxels = matr1['voxel_matrix']
    matr2 = scio.loadmat(base_path + 'hard_ke_matrix' + data_type + '.mat')
    hard_element_stiffness_train = matr2['hard_ke_matrix' + data_type]
    
    matr3 = scio.loadmat(base_path + 'hard_fe_matrix' + data_type + '.mat')
    hard_element_force_train = matr3['hard_fe_matrix' + data_type]
    matr4 = scio.loadmat(base_path + 'soft_ke_matrix' + data_type + '.mat')
    soft_element_stiffness_train = matr4['soft_ke_matrix' + data_type]
    matr5 = scio.loadmat(base_path + 'soft_fe_matrix' + data_type + '.mat')
    soft_element_force_train = matr5['soft_fe_matrix' + data_type]
    matr6 = scio.loadmat(base_path + 'C_hard' + data_type + '.mat')
    C_hard = matr6['C_hard_m' + data_type]
    matr7 = scio.loadmat(base_path + 'C_soft' + data_type + '.mat')
    C_soft = matr7['C_soft_m'+ data_type]
    matr8 = scio.loadmat(base_path + 'jacobi_matrix' + data_type + '.mat')
    jacobi = matr8['jacobi_matrix' + data_type]

    dataset = DataSet(voxels,  hard_element_stiffness_train, hard_element_force_train, soft_element_stiffness_train, soft_element_force_train, C_hard, C_soft, jacobi, voxel_length, case)
    
    # make 80% for train_dateset and 20% for test_dataet 
    # print(len(dataset))
    train_size = int(len(dataset) * 0.8)
    test_size = len(dataset) - train_size
    train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size])
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size = batch_size, sampler = DistributedSampler(train_dataset))
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size = batch_size, sampler = DistributedSampler(test_dataset))

    # train_loader = torch.utils.data.DataLoader(train_dataset, batch_size = batch_size, shuffle = shuffle)
    # test_loader = torch.utils.data.DataLoader(test_dataset, batch_size = batch_size, shuffle = shuffle)


    return train_loader, test_loader

def getDataLoader(base_path, data_type, voxel_length, case, batch_size = 8, shuffle = True):


    matr1 = scio.loadmat(base_path + 'voxel_matrix.mat')
    voxels = matr1['voxel_matrix']
    matr2 = scio.loadmat(base_path + 'hard_ke_matrix' + data_type + '.mat')
    hard_element_stiffness_train = matr2['hard_ke_matrix' + data_type]
    matr3 = scio.loadmat(base_path + 'hard_fe_matrix' + data_type + '.mat')
    hard_element_force_train = matr3['hard_fe_matrix' + data_type]
    matr4 = scio.loadmat(base_path + 'soft_ke_matrix' + data_type + '.mat')
    soft_element_stiffness_train = matr4['soft_ke_matrix' + data_type]
    matr5 = scio.loadmat(base_path + 'soft_fe_matrix' + data_type + '.mat')
    soft_element_force_train = matr5['soft_fe_matrix' + data_type]
    matr6 = scio.loadmat(base_path + 'C_hard' + data_type + '.mat')
    C_hard = matr6['C_hard_m' + data_type]
    matr7 = scio.loadmat(base_path + 'C_soft' + data_type + '.mat')
    C_soft = matr7['C_soft_m'+ data_type]
    matr8 = scio.loadmat(base_path + 'jacobi_matrix' + data_type + '.mat')
    jacobi = matr8['jacobi_matrix' + data_type]

    dataset = DataSet(voxels,  hard_element_stiffness_train, hard_element_force_train, soft_element_stiffness_train, soft_element_force_train, C_hard, C_soft, jacobi, voxel_length, case)
    
    # make 80% for train_dateset and 20% for test_dataet 
    # print(len(dataset))
    # train_size = int(len(dataset) * 0.8)
    # test_size = len(dataset) - train_size
    # train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size])
    train_loader = torch.utils.data.DataLoader(dataset, batch_size = batch_size, shuffle = shuffle, drop_last=True)
    # test_loader = torch.utils.data.DataLoader(test_dataset, batch_size = batch_size, shuffle = shuffle, drop_last=True)


    return train_loader, None